﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Reflection;
using System.Reflection.Emit;
using msilgc.cbuilder.MSILSyntaxTree;

namespace msilgc.cbuilder.MSILBasicBlocksTranslator {
    internal class BasicBlocksTranslator {
        public IEnumerable<BasicBlock> Translate(IEnumerable<Instruction> instructions, MethodBody methodBody) {
            var blocks = new List<BasicBlock>();
            var jumpTargets = FindJumpTargets(instructions).ToList();
            var currentBlock = new List<Instruction>();
            foreach (var inst in instructions) {
                if (inst.Inst == OpCodes.Ret && instructions.IndexOf(inst) == instructions.Count() - 1) {
                    if (currentBlock.Count > 0) {
                        blocks.Add(new BasicBlock(currentBlock));
                        currentBlock = new List<Instruction>();
                    }
                    var block = new BasicBlock(currentBlock);
                    block.BranchInstruction = inst;
                    blocks.Add(block);
                    currentBlock = new List<Instruction>();
                } else if (jumpTargets.Contains(inst)) {
                    if (currentBlock.Count > 0) {
                        blocks.Add(new BasicBlock(currentBlock));
                    }
                    if (inst.Inst.FlowControl == FlowControl.Branch ||
                        inst.Inst.FlowControl == FlowControl.Throw ||
                        inst.Inst.FlowControl == FlowControl.Cond_Branch) {
                        blocks.Add(new BasicBlock(new Instruction[0]) { BranchInstruction = inst });
                        currentBlock = new List<Instruction>();
                    } else {
                        currentBlock = new List<Instruction> { inst };
                    }
                } else if (IsBranch(inst)) {
                    var block = new BasicBlock(currentBlock);
                    blocks.Add(block);
                    currentBlock = new List<Instruction>();
                    block.BranchInstruction = inst;
                } else {
                    currentBlock.Add(inst);
                }
            }
            if (currentBlock.Count > 0) {
                throw new Exception("basic block not empty as expected");
            }
            foreach (var block in blocks) {
                block.Targets.AddRange(FindTargets(block, blocks, instructions));
            }
            SetSwitchInstructionTargets(instructions, blocks);
            foreach (var block in blocks) {
                ClassifyOneBlock(block, methodBody);
            }
            ValidateOnlyBlocksOfTheSameTypeCanReachEachother(blocks);
            ValidateBlocksDontContainAnyBranchInstructions(blocks);
            return blocks;
        }

        private static void ValidateBlocksDontContainAnyBranchInstructions(IEnumerable<BasicBlock> blocks) {
            foreach (var block in blocks) {
                foreach (var instruction in block.Instructions) {
                    if (instruction.Inst.FlowControl == FlowControl.Branch ||
                        instruction.Inst.FlowControl == FlowControl.Break ||
                        instruction.Inst.FlowControl == FlowControl.Return ||
                        instruction.Inst.FlowControl == FlowControl.Throw ||
                        instruction.Inst.FlowControl == FlowControl.Cond_Branch) {
                        throw new Exception("Basic block is invalid - has a branch in the body - " + instruction.Inst);
                    }
                }
            }
        }

        private static void SetSwitchInstructionTargets(IEnumerable<Instruction> instructions, List<BasicBlock> blocks) {
            foreach (var instruction in instructions) {
                if (instruction.SwitchOffsets != null) {
                    instruction.SwitchTargets = new IBasicBlock[instruction.SwitchOffsets.Length];
                    var nextInst = instructions.ElementAt(instructions.IndexOf(instruction) + 1);
                    for (var i = 0; i < instruction.SwitchOffsets.Length; i++) {
                        var targetAddress = nextInst.Offset + instruction.SwitchOffsets[i];
                        var targetBlock = blocks.First(block => block.Offset == targetAddress);
                        instruction.SwitchTargets[i] = targetBlock;
                    }
                }
            }
        }

        private static void ValidateOnlyBlocksOfTheSameTypeCanReachEachother(IEnumerable<BasicBlock> blocks) {
            foreach (var block in blocks) {
                var reachable = FindReachableBlocks(block);
                if (reachable.Any(target =>
                                  block.IsNormal != target.IsNormal || block.IsFinally != target.IsFinally ||
                                  block.IsExceptionHandler != target.IsExceptionHandler)) {
                    throw new Exception("Blocks of two different types should not be able to reach eachother");
                }
            }
        }

        private static IEnumerable<BasicBlock> FindReachableBlocks(BasicBlock block) {
            var reached = new List<BasicBlock> {block};
            var count = 0;
            while (count != reached.Count) {
                count = reached.Count;
                var targetSet =
                    from b in reached
                    from target in b.Targets
                    where !reached.Contains(target)
                    select target;
                reached.AddRange(targetSet.ToArray());
            }
            return reached;
        }

        private static void ClassifyOneBlock(BasicBlock block, MethodBody methodBody) {
            foreach (var clause in methodBody.ExceptionHandlingClauses) {
                if (clause.HandlerOffset <= block.Offset && (clause.HandlerOffset + clause.HandlerLength) > block.Offset) {
                    if (clause.Flags == ExceptionHandlingClauseOptions.Clause) {
                        block.IsExceptionHandler = true;
                        block.ExceptionTypeIfExceptionHandler = clause.CatchType;
                        return;
                    } else if ((clause.Flags & ExceptionHandlingClauseOptions.Finally) != 0) {
                        block.IsFinally = true;
                        return;
                    }
                } 
            }
        }

        private static IEnumerable<BasicBlock> FindTargets(BasicBlock from, IEnumerable<BasicBlock> possibleTargets, IEnumerable<Instruction> instructions) {
            if (from.BranchInstruction == null) {
                yield return possibleTargets.ElementAt(possibleTargets.IndexOf(from) + 1);
            } else if (from.BranchInstruction.Inst.FlowControl == FlowControl.Cond_Branch) {
                var nextInst = instructions.ElementAt(instructions.IndexOf(from.BranchInstruction) + 1);
                var target = TranslateTarget(@from.BranchInstruction.Data, nextInst.Offset);
                yield return possibleTargets.First(block => block.Offset == target);
                var nextBlock = possibleTargets.ElementAt(possibleTargets.IndexOf(from) + 1);
                yield return nextBlock;
            } else if (from.BranchInstruction.Inst.FlowControl == FlowControl.Branch) {
                var nextInst = instructions.ElementAt(instructions.IndexOf(from.BranchInstruction) + 1);
                var target = TranslateTarget(@from.BranchInstruction.Data, nextInst.Offset);
                yield return possibleTargets.First(block => block.Offset == target);
            }
        }

        private static int TranslateTarget(byte[] data, int offset) {
            if (data.Length == 1) {
                unchecked {
                    sbyte target = (sbyte)data[0];
                    return offset + target;
                }
            } else if (data.Length == 4) {
                int target = (
                                 ((int)data[3] << 3 * 8) &
                                 ((int)data[2] << 2 * 8) &
                                 ((int)data[1] << 1 * 8) &
                                 ((int)data[0] << 0 * 8)
                             );
                return offset + target;
            } else {
                throw new Exception("did not expect data length == " + data.Length);
            }
        }

        private static bool IsBranch(Instruction inst) {
            return
                inst.Inst.FlowControl == FlowControl.Branch ||
                inst.Inst.FlowControl == FlowControl.Cond_Branch ||
                inst.Inst.FlowControl == FlowControl.Return;
        }

        private static IEnumerable<Instruction> FindJumpInstructions(IEnumerable<Instruction> instructions) {
            return instructions.Where(IsBranch);
        }

        private static IEnumerable<Instruction> FindJumpTargets(IEnumerable<Instruction> instructions) {
            foreach (var jumpInst in FindJumpInstructions(instructions)) {
                if (jumpInst.Inst.FlowControl != FlowControl.Return) {
                    var nextInst = instructions.ElementAt(instructions.IndexOf(jumpInst) + 1);
                    int target = TranslateTarget(jumpInst.Data, nextInst.Offset);
                    yield return instructions.First(inst => inst.Offset == target);
                }
            }
        }
    }
}